package it.uniroma3.dia.ia.textcategorization.categorizer.opennlp;

import java.io.IOException;
import java.util.Map;
import java.util.Set;
import it.uniroma3.dia.ia.textcategorization.categorizer.Categorizer;
import it.uniroma3.dia.ia.textcategorization.categorizer.CategorizerBuilder;
import it.uniroma3.dia.ia.textcategorization.categorizer.CategorizerBuilderException;
import it.uniroma3.dia.ia.textcategorization.model.Category;
import it.uniroma3.dia.ia.textcategorization.model.Text;
import opennlp.tools.doccat.DoccatModel;
import opennlp.tools.doccat.DocumentCategorizerME;
import opennlp.tools.doccat.DocumentSample;
import opennlp.tools.util.ObjectStream;

public final class OpenNLPCategorizerBuilder implements CategorizerBuilder {

	@Override
	public Categorizer newCategorizer(Set<Category> categories,
			Map<Text, Category> trainingSet) {
		try {
			ObjectStream<DocumentSample> trainingMaterial = OpenNLPTrainingFormat
					.getFormattedTrainingMaterial(trainingSet);
			DoccatModel model = DocumentCategorizerME.train("en",
					trainingMaterial);
			DocumentCategorizerME categorizer = new DocumentCategorizerME(model);
			return new OpenNLPCategorizer(categories, categorizer);
		} catch (IOException exception) {
			throw new CategorizerBuilderException(
					"Exception during categorizer's training", exception);
		}
	}

}
